{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.chdir('../')\n",
    "os.environ['CUDA_VISIBLE_DEVICES']='0'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, re, pickle\n",
    "\n",
    "import tensorflow as tf\n",
    "import tf_extend as tfe\n",
    "\n",
    "from nets.mtcnn import MTCNN\n",
    "\n",
    "mtcnn_obj = MTCNN()\n",
    "\n",
    "def get_annotations():\n",
    "    anno_file = open('/home/luojiapeng/datasets/widerface/wider_face_split/wider_face_val_bbx_gt.txt', 'r')\n",
    "    annotations = []\n",
    "    for i, line in enumerate(anno_file.readlines()):\n",
    "        line = line.rstrip()\n",
    "        if re.match('^.*\\.jpg$', line):\n",
    "            annotations.append(line)\n",
    "        elif len(line.split()) > 4:\n",
    "            info = ' '.join(line.split()[:4])\n",
    "            annotations[-1] = annotations[-1] + ' ' + info\n",
    "    return annotations\n",
    "\n",
    "\n",
    "def get_dataset():\n",
    "    base_dir = '/home/luojiapeng/datasets/widerface/WIDER_val/images'\n",
    "    anno = get_annotations()\n",
    "    fnames = [x.split()[0] for x in anno]\n",
    "    fnames = [os.path.join(base_dir, x) for x in fnames]\n",
    "    dataset = tf.data.Dataset.from_tensor_slices((fnames,))\n",
    "\n",
    "    def fn(fname):\n",
    "        raw_img = tf.read_file(fname)\n",
    "        image = tf.image.decode_jpeg(raw_img, channels=3)\n",
    "        image = tf.cast(image, tf.float32)\n",
    "        return {'fname': fname, 'image': image}\n",
    "\n",
    "    dataset = dataset.map(fn, num_parallel_calls=8)\n",
    "    return dataset\n",
    "\n",
    "def save_wider_result(output_dir, fnames, result):\n",
    "    if not os.path.exists(output_dir):\n",
    "        os.mkdir(output_dir)\n",
    "    current_event = ''\n",
    "    for i in range(len(fnames)):\n",
    "        fname = fnames[i]\n",
    "        bboxes, scores = result[i]['bboxes'], result[i]['scores']\n",
    "        assert len(bboxes) == len(scores)\n",
    "        event = fname.split('/')[0]\n",
    "        if current_event != event:\n",
    "            current_event = event\n",
    "            save_path = os.path.join(output_dir, current_event)\n",
    "            if not os.path.exists(save_path):\n",
    "                os.mkdir(save_path)\n",
    "            print('current path:', current_event)\n",
    "\n",
    "        out_fname = fname.split('.jpg')[0]\n",
    "        out_fname = os.path.join(output_dir, out_fname + '.txt')\n",
    "        fid = open(out_fname, 'w')\n",
    "        fid.write(fname.split('/')[-1] + '\\n')\n",
    "        if bboxes is None:\n",
    "            fid.write(str(1) + '\\n')\n",
    "            fid.write('%f %f %f %f %f\\n' % (0, 0, 0, 0, 0.99))\n",
    "            continue\n",
    "        else:\n",
    "            fid.write(str(len(bboxes)) + '\\n')\n",
    "            for _i in range(len(scores)):\n",
    "                s, b =scores[_i], bboxes[_i]\n",
    "                fid.write('%.2f %.2f %.2f %.2f %.2f\\n' % (b[1], b[0], b[3] - b[1] + 1, b[2] - b[0] + 1, s))\n",
    "\n",
    "            fid.close()\n",
    "            if i % 100 == 0 and i:\n",
    "                print(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:tensorflow:get_checkpoint_dir: result/pnet_mse_t3_0_rnet_l1t3_0_onet_pin_0_0\n",
      "INFO:tensorflow:Using config: {'_model_dir': 'checkpoints/pnet_mse_t3_0', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': 600, '_session_config': None, '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_train_distribute': None, '_device_fn': None, '_service': None, '_cluster_spec': <tensorflow.python.training.server_lib.ClusterSpec object at 0x7f3a8c25c390>, '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1}\n",
      "INFO:tensorflow:Calling model_fn.\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "INFO:tensorflow:Graph was finalized.\n",
      "INFO:tensorflow:Restoring parameters from checkpoints/pnet_mse_t3_0/model.ckpt-200000\n",
      "INFO:tensorflow:Running local_init_op.\n",
      "INFO:tensorflow:Done running local_init_op.\n",
      "INFO:tensorflow:Using config: {'_model_dir': 'checkpoints/rnet_l1t3_0', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': 600, '_session_config': None, '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_train_distribute': None, '_device_fn': None, '_service': None, '_cluster_spec': <tensorflow.python.training.server_lib.ClusterSpec object at 0x7f3a8c260da0>, '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1}\n",
      "INFO:tensorflow:Calling model_fn.\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "INFO:tensorflow:Graph was finalized.\n",
      "INFO:tensorflow:Restoring parameters from checkpoints/rnet_l1t3_0/model.ckpt-200000\n",
      "INFO:tensorflow:Running local_init_op.\n",
      "INFO:tensorflow:Done running local_init_op.\n",
      "INFO:tensorflow:\n",
      "INFO:tensorflow: (0.674 sec)\n",
      "INFO:tensorflow: (0.472 sec)\n",
      "INFO:tensorflow: (0.407 sec)\n",
      "INFO:tensorflow: (0.390 sec)\n",
      "INFO:tensorflow: (0.413 sec)\n",
      "INFO:tensorflow: (0.357 sec)\n",
      "INFO:tensorflow: (0.408 sec)\n",
      "INFO:tensorflow: (0.355 sec)\n",
      "INFO:tensorflow: (0.379 sec)\n",
      "INFO:tensorflow: (0.377 sec)\n",
      "INFO:tensorflow: (0.392 sec)\n",
      "INFO:tensorflow: (0.335 sec)\n",
      "INFO:tensorflow: (0.373 sec)\n",
      "INFO:tensorflow: (0.328 sec)\n",
      "INFO:tensorflow: (0.353 sec)\n",
      "INFO:tensorflow: (0.383 sec)\n",
      "INFO:tensorflow: (0.358 sec)\n",
      "INFO:tensorflow: (0.341 sec)\n",
      "INFO:tensorflow: (0.363 sec)\n",
      "INFO:tensorflow: (0.350 sec)\n",
      "INFO:tensorflow: (0.341 sec)\n",
      "INFO:tensorflow: (0.330 sec)\n",
      "INFO:tensorflow: (0.338 sec)\n",
      "INFO:tensorflow: (0.345 sec)\n",
      "INFO:tensorflow: (0.366 sec)\n",
      "INFO:tensorflow: (0.342 sec)\n",
      "INFO:tensorflow: (0.336 sec)\n",
      "INFO:tensorflow: (0.352 sec)\n",
      "INFO:tensorflow: (0.383 sec)\n",
      "INFO:tensorflow: (0.367 sec)\n",
      "INFO:tensorflow: (0.408 sec)\n",
      "INFO:tensorflow: (0.372 sec)\n",
      "INFO:tensorflow:Using config: {'_model_dir': 'checkpoints/onet_pin_0', '_tf_random_seed': None, '_save_summary_steps': 100, '_save_checkpoints_steps': None, '_save_checkpoints_secs': 600, '_session_config': None, '_keep_checkpoint_max': 5, '_keep_checkpoint_every_n_hours': 10000, '_log_step_count_steps': 100, '_train_distribute': None, '_device_fn': None, '_service': None, '_cluster_spec': <tensorflow.python.training.server_lib.ClusterSpec object at 0x7f39a025f438>, '_task_type': 'worker', '_task_id': 0, '_global_id_in_cluster': 0, '_master': '', '_evaluation_master': '', '_is_chief': True, '_num_ps_replicas': 0, '_num_worker_replicas': 1}\n",
      "INFO:tensorflow:Calling model_fn.\n",
      "INFO:tensorflow:Done calling model_fn.\n",
      "INFO:tensorflow:Graph was finalized.\n",
      "INFO:tensorflow:Restoring parameters from checkpoints/onet_pin_0/model.ckpt-300000\n",
      "INFO:tensorflow:Running local_init_op.\n",
      "INFO:tensorflow:Done running local_init_op.\n",
      "INFO:tensorflow:\n",
      "INFO:tensorflow: (0.644 sec)\n",
      "INFO:tensorflow: (0.478 sec)\n",
      "INFO:tensorflow: (0.453 sec)\n",
      "INFO:tensorflow: (0.427 sec)\n",
      "INFO:tensorflow: (0.377 sec)\n",
      "INFO:tensorflow: (0.338 sec)\n",
      "INFO:tensorflow: (0.378 sec)\n",
      "INFO:tensorflow: (0.434 sec)\n",
      "INFO:tensorflow: (0.425 sec)\n",
      "INFO:tensorflow: (0.430 sec)\n",
      "INFO:tensorflow: (0.411 sec)\n",
      "INFO:tensorflow: (0.381 sec)\n",
      "INFO:tensorflow: (0.376 sec)\n",
      "INFO:tensorflow: (0.399 sec)\n",
      "INFO:tensorflow: (0.400 sec)\n",
      "INFO:tensorflow: (0.386 sec)\n",
      "INFO:tensorflow: (0.405 sec)\n",
      "INFO:tensorflow: (0.371 sec)\n",
      "INFO:tensorflow: (0.407 sec)\n",
      "INFO:tensorflow: (0.350 sec)\n",
      "INFO:tensorflow: (0.367 sec)\n",
      "INFO:tensorflow: (0.379 sec)\n",
      "INFO:tensorflow: (0.383 sec)\n",
      "INFO:tensorflow: (0.394 sec)\n",
      "INFO:tensorflow: (0.366 sec)\n",
      "INFO:tensorflow: (0.330 sec)\n",
      "INFO:tensorflow: (0.352 sec)\n",
      "INFO:tensorflow: (0.390 sec)\n",
      "INFO:tensorflow: (0.367 sec)\n",
      "INFO:tensorflow: (0.390 sec)\n",
      "INFO:tensorflow: (0.400 sec)\n",
      "INFO:tensorflow: (0.405 sec)\n",
      "current path: 0--Parade\n",
      "100\n",
      "current path: 1--Handshaking\n",
      "current path: 10--People_Marching\n",
      "200\n",
      "current path: 11--Meeting\n",
      "current path: 12--Group\n",
      "300\n",
      "current path: 13--Interview\n",
      "400\n",
      "500\n",
      "current path: 14--Traffic\n",
      "current path: 15--Stock_Market\n",
      "current path: 16--Award_Ceremony\n",
      "600\n",
      "current path: 17--Ceremony\n",
      "current path: 18--Concerts\n",
      "current path: 19--Couple\n",
      "700\n",
      "current path: 2--Demonstration\n",
      "800\n",
      "900\n",
      "current path: 20--Family_Group\n",
      "1000\n",
      "current path: 21--Festival\n",
      "current path: 22--Picnic\n",
      "1100\n",
      "current path: 23--Shoppers\n",
      "current path: 24--Soldier_Firing\n",
      "current path: 25--Soldier_Patrol\n",
      "1200\n",
      "current path: 26--Soldier_Drilling\n",
      "current path: 27--Spa\n",
      "current path: 28--Sports_Fan\n",
      "1300\n",
      "current path: 29--Students_Schoolkids\n",
      "current path: 3--Riot\n",
      "1400\n",
      "current path: 30--Surgeons\n",
      "current path: 31--Waiter_Waitress\n",
      "1500\n",
      "current path: 32--Worker_Laborer\n",
      "current path: 33--Running\n",
      "current path: 34--Baseball\n",
      "1600\n",
      "current path: 35--Basketball\n",
      "1700\n",
      "current path: 36--Football\n",
      "current path: 37--Soccer\n",
      "1800\n",
      "current path: 38--Tennis\n",
      "current path: 39--Ice_Skating\n",
      "1900\n",
      "current path: 4--Dancing\n",
      "current path: 40--Gymnastics\n",
      "2000\n",
      "current path: 41--Swimming\n",
      "2100\n",
      "current path: 42--Car_Racing\n",
      "current path: 43--Row_Boat\n",
      "2200\n",
      "current path: 44--Aerobics\n",
      "current path: 45--Balloonist\n",
      "current path: 46--Jockey\n",
      "2300\n",
      "current path: 47--Matador_Bullfighter\n",
      "current path: 48--Parachutist_Paratrooper\n",
      "2400\n",
      "current path: 49--Greeting\n",
      "current path: 5--Car_Accident\n",
      "current path: 50--Celebration_Or_Party\n",
      "2500\n",
      "current path: 51--Dresses\n",
      "2600\n",
      "current path: 52--Photographers\n",
      "current path: 53--Raid\n",
      "2700\n",
      "current path: 54--Rescue\n",
      "current path: 55--Sports_Coach_Trainer\n",
      "2800\n",
      "current path: 56--Voter\n",
      "current path: 57--Angler\n",
      "2900\n",
      "current path: 58--Hockey\n",
      "current path: 59--people--driving--car\n",
      "3000\n",
      "current path: 6--Funeral\n",
      "current path: 61--Street_Battle\n",
      "current path: 7--Cheering\n",
      "3100\n",
      "current path: 8--Election_Campain\n",
      "current path: 9--Press_Conference\n",
      "3200\n",
      "current path: 0--Parade\n",
      "100\n",
      "current path: 1--Handshaking\n",
      "current path: 10--People_Marching\n",
      "200\n",
      "current path: 11--Meeting\n",
      "current path: 12--Group\n",
      "300\n",
      "current path: 13--Interview\n",
      "400\n",
      "500\n",
      "current path: 14--Traffic\n",
      "current path: 15--Stock_Market\n",
      "current path: 16--Award_Ceremony\n",
      "600\n",
      "current path: 17--Ceremony\n",
      "current path: 18--Concerts\n",
      "current path: 19--Couple\n",
      "700\n",
      "current path: 2--Demonstration\n",
      "800\n",
      "900\n",
      "current path: 20--Family_Group\n",
      "1000\n",
      "current path: 21--Festival\n",
      "current path: 22--Picnic\n",
      "1100\n",
      "current path: 23--Shoppers\n",
      "current path: 24--Soldier_Firing\n",
      "current path: 25--Soldier_Patrol\n",
      "1200\n",
      "current path: 26--Soldier_Drilling\n",
      "current path: 27--Spa\n",
      "current path: 28--Sports_Fan\n",
      "1300\n",
      "current path: 29--Students_Schoolkids\n",
      "current path: 3--Riot\n",
      "1400\n",
      "current path: 30--Surgeons\n",
      "current path: 31--Waiter_Waitress\n",
      "1500\n",
      "current path: 32--Worker_Laborer\n",
      "current path: 33--Running\n",
      "current path: 34--Baseball\n",
      "1600\n",
      "current path: 35--Basketball\n",
      "1700\n",
      "current path: 36--Football\n",
      "current path: 37--Soccer\n",
      "1800\n",
      "current path: 38--Tennis\n",
      "current path: 39--Ice_Skating\n",
      "1900\n",
      "current path: 4--Dancing\n",
      "current path: 40--Gymnastics\n",
      "2000\n",
      "current path: 41--Swimming\n",
      "2100\n",
      "current path: 42--Car_Racing\n",
      "current path: 43--Row_Boat\n",
      "2200\n",
      "current path: 44--Aerobics\n",
      "current path: 45--Balloonist\n",
      "current path: 46--Jockey\n",
      "2300\n",
      "current path: 47--Matador_Bullfighter\n",
      "current path: 48--Parachutist_Paratrooper\n",
      "2400\n",
      "current path: 49--Greeting\n",
      "current path: 5--Car_Accident\n",
      "current path: 50--Celebration_Or_Party\n",
      "2500\n",
      "current path: 51--Dresses\n",
      "2600\n",
      "current path: 52--Photographers\n",
      "current path: 53--Raid\n",
      "2700\n",
      "current path: 54--Rescue\n",
      "current path: 55--Sports_Coach_Trainer\n",
      "2800\n",
      "current path: 56--Voter\n",
      "current path: 57--Angler\n",
      "2900\n",
      "current path: 58--Hockey\n",
      "current path: 59--people--driving--car\n",
      "3000\n",
      "current path: 6--Funeral\n",
      "current path: 61--Street_Battle\n",
      "current path: 7--Cheering\n",
      "3100\n",
      "current path: 8--Election_Campain\n",
      "current path: 9--Press_Conference\n",
      "3200\n",
      "current path: 0--Parade\n",
      "100\n",
      "current path: 1--Handshaking\n",
      "current path: 10--People_Marching\n",
      "200\n",
      "current path: 11--Meeting\n",
      "current path: 12--Group\n",
      "300\n",
      "current path: 13--Interview\n",
      "400\n",
      "500\n",
      "current path: 14--Traffic\n",
      "current path: 15--Stock_Market\n",
      "current path: 16--Award_Ceremony\n",
      "600\n",
      "current path: 17--Ceremony\n",
      "current path: 18--Concerts\n",
      "current path: 19--Couple\n",
      "700\n",
      "current path: 2--Demonstration\n",
      "800\n",
      "900\n",
      "current path: 20--Family_Group\n",
      "1000\n",
      "current path: 21--Festival\n",
      "current path: 22--Picnic\n",
      "1100\n",
      "current path: 23--Shoppers\n",
      "current path: 24--Soldier_Firing\n",
      "current path: 25--Soldier_Patrol\n",
      "1200\n",
      "current path: 26--Soldier_Drilling\n",
      "current path: 27--Spa\n",
      "current path: 28--Sports_Fan\n",
      "1300\n",
      "current path: 29--Students_Schoolkids\n",
      "current path: 3--Riot\n",
      "1400\n",
      "current path: 30--Surgeons\n",
      "current path: 31--Waiter_Waitress\n",
      "1500\n",
      "current path: 32--Worker_Laborer\n",
      "current path: 33--Running\n",
      "current path: 34--Baseball\n",
      "1600\n",
      "current path: 35--Basketball\n",
      "1700\n",
      "current path: 36--Football\n",
      "current path: 37--Soccer\n",
      "1800\n",
      "current path: 38--Tennis\n",
      "current path: 39--Ice_Skating\n",
      "1900\n",
      "current path: 4--Dancing\n",
      "current path: 40--Gymnastics\n",
      "2000\n",
      "current path: 41--Swimming\n",
      "2100\n",
      "current path: 42--Car_Racing\n",
      "current path: 43--Row_Boat\n",
      "2200\n",
      "current path: 44--Aerobics\n",
      "current path: 45--Balloonist\n",
      "current path: 46--Jockey\n",
      "2300\n",
      "current path: 47--Matador_Bullfighter\n",
      "current path: 48--Parachutist_Paratrooper\n",
      "2400\n",
      "current path: 49--Greeting\n",
      "current path: 5--Car_Accident\n",
      "current path: 50--Celebration_Or_Party\n",
      "2500\n",
      "current path: 51--Dresses\n",
      "2600\n",
      "current path: 52--Photographers\n",
      "current path: 53--Raid\n",
      "2700\n",
      "current path: 54--Rescue\n",
      "current path: 55--Sports_Coach_Trainer\n",
      "2800\n",
      "current path: 56--Voter\n",
      "current path: 57--Angler\n",
      "2900\n",
      "current path: 58--Hockey\n",
      "current path: 59--people--driving--car\n",
      "3000\n",
      "current path: 6--Funeral\n",
      "current path: 61--Street_Battle\n",
      "current path: 7--Cheering\n",
      "3100\n",
      "current path: 8--Election_Campain\n",
      "current path: 9--Press_Conference\n",
      "3200\n"
     ]
    }
   ],
   "source": [
    "output_dir = tfe.get_checkpoint_dir('result', 'pnet_mse_t3_0_rnet_l1t3_0_onet_pin_0')\n",
    "if not os.path.exists(output_dir):\n",
    "    os.mkdir(output_dir)\n",
    "params = {\n",
    "    'pnet_model_dir': 'checkpoints/pnet_mse_t3_0',\n",
    "    'pnet_thres': 0.3,\n",
    "    'rnet_model_dir': 'checkpoints/rnet_l1t3_0',\n",
    "    'rnet_thres': 0.3,\n",
    "    'onet_model_dir': 'checkpoints/onet_pin_0',\n",
    "    'onet_thres': 0.3\n",
    "}\n",
    "p_res, r_res, o_res = mtcnn_obj.predict(get_dataset, params)\n",
    "fnames = [x.split()[0] for x in get_annotations()]\n",
    "save_wider_result(os.path.join(output_dir, 'pnet'), fnames, p_res)\n",
    "save_wider_result(os.path.join(output_dir, 'rnet'), fnames, r_res)\n",
    "save_wider_result(os.path.join(output_dir, 'onet'), fnames, o_res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
